{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OpenAI Assistant Agent\n",
    "\n",
    "[Open AI Assistant](https://platform.openai.com/docs/assistants/overview) \n",
    "and [Azure OpenAI Assistant](https://learn.microsoft.com/en-us/azure/ai-services/openai/how-to/assistant)\n",
    "are server-side APIs for building\n",
    "agents.\n",
    "They can be used to build agents in AutoGen. This cookbook demonstrates how to\n",
    "to use OpenAI Assistant to create an agent that can run code and Q&A over document.\n",
    "\n",
    "## Message Protocol\n",
    "\n",
    "First, we need to specify the message protocol for the agent backed by \n",
    "OpenAI Assistant. The message protocol defines the structure of messages\n",
    "handled and published by the agent. \n",
    "For illustration, we define a simple\n",
    "message protocol of 4 message types: `Message`, `Reset`, `UploadForCodeInterpreter` and `UploadForFileSearch`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class TextMessage:\n",
    "    content: str\n",
    "    source: str\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Reset:\n",
    "    pass\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class UploadForCodeInterpreter:\n",
    "    file_path: str\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class UploadForFileSearch:\n",
    "    file_path: str\n",
    "    vector_store_id: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `TextMessage` message type is used to communicate with the agent. It has a\n",
    "`content` field that contains the message content, and a `source` field\n",
    "for the sender. The `Reset` message type is a control message that resets\n",
    "the memory of the agent. It has no fields. This is useful when we need to\n",
    "start a new conversation with the agent.\n",
    "\n",
    "The `UploadForCodeInterpreter` message type is used to upload data files\n",
    "for the code interpreter and `UploadForFileSearch` message type is used to upload\n",
    "documents for file search. Both message types have a `file_path` field that contains\n",
    "the local path to the file to be uploaded.\n",
    "\n",
    "## Defining the Agent\n",
    "\n",
    "Next, we define the agent class.\n",
    "The agent class constructor has the following arguments: `description`,\n",
    "`client`, `assistant_id`, `thread_id`, and `assistant_event_handler_factory`.\n",
    "The `client` argument is the OpenAI async client object, and the\n",
    "`assistant_event_handler_factory` is for creating an assistant event handler\n",
    "to handle OpenAI Assistant events.\n",
    "This can be used to create streaming output from the assistant.\n",
    "\n",
    "The agent class has the following message handlers:\n",
    "- `handle_message`: Handles the `TextMessage` message type, and sends back the\n",
    "  response from the assistant.\n",
    "- `handle_reset`: Handles the `Reset` message type, and resets the memory\n",
    "    of the assistant agent.\n",
    "- `handle_upload_for_code_interpreter`: Handles the `UploadForCodeInterpreter`\n",
    "  message type, and uploads the file to the code interpreter.\n",
    "- `handle_upload_for_file_search`: Handles the `UploadForFileSearch`\n",
    "    message type, and uploads the document to the file search.\n",
    "\n",
    "\n",
    "The memory of the assistant is stored inside a thread, which is kept in the\n",
    "server side. The thread is referenced by the `thread_id` argument."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "import os\n",
    "from typing import Any, Callable, List\n",
    "\n",
    "import aiofiles\n",
    "from autogen_core import AgentId, MessageContext, RoutedAgent, message_handler\n",
    "from openai import AsyncAssistantEventHandler, AsyncClient\n",
    "from openai.types.beta.thread import ToolResources, ToolResourcesFileSearch\n",
    "\n",
    "\n",
    "class OpenAIAssistantAgent(RoutedAgent):\n",
    "    \"\"\"An agent implementation that uses the OpenAI Assistant API to generate\n",
    "    responses.\n",
    "\n",
    "    Args:\n",
    "        description (str): The description of the agent.\n",
    "        client (openai.AsyncClient): The client to use for the OpenAI API.\n",
    "        assistant_id (str): The assistant ID to use for the OpenAI API.\n",
    "        thread_id (str): The thread ID to use for the OpenAI API.\n",
    "        assistant_event_handler_factory (Callable[[], AsyncAssistantEventHandler], optional):\n",
    "            A factory function to create an async assistant event handler. Defaults to None.\n",
    "            If provided, the agent will use the streaming mode with the event handler.\n",
    "            If not provided, the agent will use the blocking mode to generate responses.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        description: str,\n",
    "        client: AsyncClient,\n",
    "        assistant_id: str,\n",
    "        thread_id: str,\n",
    "        assistant_event_handler_factory: Callable[[], AsyncAssistantEventHandler],\n",
    "    ) -> None:\n",
    "        super().__init__(description)\n",
    "        self._client = client\n",
    "        self._assistant_id = assistant_id\n",
    "        self._thread_id = thread_id\n",
    "        self._assistant_event_handler_factory = assistant_event_handler_factory\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_message(self, message: TextMessage, ctx: MessageContext) -> TextMessage:\n",
    "        \"\"\"Handle a message. This method adds the message to the thread and publishes a response.\"\"\"\n",
    "        # Save the message to the thread.\n",
    "        await ctx.cancellation_token.link_future(\n",
    "            asyncio.ensure_future(\n",
    "                self._client.beta.threads.messages.create(\n",
    "                    thread_id=self._thread_id,\n",
    "                    content=message.content,\n",
    "                    role=\"user\",\n",
    "                    metadata={\"sender\": message.source},\n",
    "                )\n",
    "            )\n",
    "        )\n",
    "        # Generate a response.\n",
    "        async with self._client.beta.threads.runs.stream(\n",
    "            thread_id=self._thread_id,\n",
    "            assistant_id=self._assistant_id,\n",
    "            event_handler=self._assistant_event_handler_factory(),\n",
    "        ) as stream:\n",
    "            await ctx.cancellation_token.link_future(asyncio.ensure_future(stream.until_done()))\n",
    "\n",
    "        # Get the last message.\n",
    "        messages = await ctx.cancellation_token.link_future(\n",
    "            asyncio.ensure_future(self._client.beta.threads.messages.list(self._thread_id, order=\"desc\", limit=1))\n",
    "        )\n",
    "        last_message_content = messages.data[0].content\n",
    "\n",
    "        # Get the text content from the last message.\n",
    "        text_content = [content for content in last_message_content if content.type == \"text\"]\n",
    "        if not text_content:\n",
    "            raise ValueError(f\"Expected text content in the last message: {last_message_content}\")\n",
    "\n",
    "        return TextMessage(content=text_content[0].text.value, source=self.metadata[\"type\"])\n",
    "\n",
    "    @message_handler()\n",
    "    async def on_reset(self, message: Reset, ctx: MessageContext) -> None:\n",
    "        \"\"\"Handle a reset message. This method deletes all messages in the thread.\"\"\"\n",
    "        # Get all messages in this thread.\n",
    "        all_msgs: List[str] = []\n",
    "        while True:\n",
    "            if not all_msgs:\n",
    "                msgs = await ctx.cancellation_token.link_future(\n",
    "                    asyncio.ensure_future(self._client.beta.threads.messages.list(self._thread_id))\n",
    "                )\n",
    "            else:\n",
    "                msgs = await ctx.cancellation_token.link_future(\n",
    "                    asyncio.ensure_future(self._client.beta.threads.messages.list(self._thread_id, after=all_msgs[-1]))\n",
    "                )\n",
    "            for msg in msgs.data:\n",
    "                all_msgs.append(msg.id)\n",
    "            if not msgs.has_next_page():\n",
    "                break\n",
    "        # Delete all the messages.\n",
    "        for msg_id in all_msgs:\n",
    "            status = await ctx.cancellation_token.link_future(\n",
    "                asyncio.ensure_future(\n",
    "                    self._client.beta.threads.messages.delete(message_id=msg_id, thread_id=self._thread_id)\n",
    "                )\n",
    "            )\n",
    "            assert status.deleted is True\n",
    "\n",
    "    @message_handler()\n",
    "    async def on_upload_for_code_interpreter(self, message: UploadForCodeInterpreter, ctx: MessageContext) -> None:\n",
    "        \"\"\"Handle an upload for code interpreter. This method uploads a file and updates the thread with the file.\"\"\"\n",
    "        # Get the file content.\n",
    "        async with aiofiles.open(message.file_path, mode=\"rb\") as f:\n",
    "            file_content = await ctx.cancellation_token.link_future(asyncio.ensure_future(f.read()))\n",
    "        file_name = os.path.basename(message.file_path)\n",
    "        # Upload the file.\n",
    "        file = await ctx.cancellation_token.link_future(\n",
    "            asyncio.ensure_future(self._client.files.create(file=(file_name, file_content), purpose=\"assistants\"))\n",
    "        )\n",
    "        # Get existing file ids from tool resources.\n",
    "        thread = await ctx.cancellation_token.link_future(\n",
    "            asyncio.ensure_future(self._client.beta.threads.retrieve(thread_id=self._thread_id))\n",
    "        )\n",
    "        tool_resources: ToolResources = thread.tool_resources if thread.tool_resources else ToolResources()\n",
    "        assert tool_resources.code_interpreter is not None\n",
    "        if tool_resources.code_interpreter.file_ids:\n",
    "            file_ids = tool_resources.code_interpreter.file_ids\n",
    "        else:\n",
    "            file_ids = [file.id]\n",
    "        # Update thread with new file.\n",
    "        await ctx.cancellation_token.link_future(\n",
    "            asyncio.ensure_future(\n",
    "                self._client.beta.threads.update(\n",
    "                    thread_id=self._thread_id,\n",
    "                    tool_resources={\n",
    "                        \"code_interpreter\": {\"file_ids\": file_ids},\n",
    "                    },\n",
    "                )\n",
    "            )\n",
    "        )\n",
    "\n",
    "    @message_handler()\n",
    "    async def on_upload_for_file_search(self, message: UploadForFileSearch, ctx: MessageContext) -> None:\n",
    "        \"\"\"Handle an upload for file search. This method uploads a file and updates the vector store.\"\"\"\n",
    "        # Get the file content.\n",
    "        async with aiofiles.open(message.file_path, mode=\"rb\") as file:\n",
    "            file_content = await ctx.cancellation_token.link_future(asyncio.ensure_future(file.read()))\n",
    "        file_name = os.path.basename(message.file_path)\n",
    "        # Upload the file.\n",
    "        await ctx.cancellation_token.link_future(\n",
    "            asyncio.ensure_future(\n",
    "                self._client.vector_stores.file_batches.upload_and_poll(\n",
    "                    vector_store_id=message.vector_store_id,\n",
    "                    files=[(file_name, file_content)],\n",
    "                )\n",
    "            )\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The agent class is a thin wrapper around the OpenAI Assistant API to implement\n",
    "the message protocol. More features, such as multi-modal message handling,\n",
    "can be added by extending the message protocol."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assistant Event Handler\n",
    "\n",
    "The assistant event handler provides call-backs for handling Assistant API\n",
    "specific events. This is useful for handling streaming output from the assistant\n",
    "and further user interface integration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openai import AsyncAssistantEventHandler, AsyncClient\n",
    "from openai.types.beta.threads import Message, Text, TextDelta\n",
    "from openai.types.beta.threads.runs import RunStep, RunStepDelta\n",
    "from typing_extensions import override\n",
    "\n",
    "\n",
    "class EventHandler(AsyncAssistantEventHandler):\n",
    "    @override\n",
    "    async def on_text_delta(self, delta: TextDelta, snapshot: Text) -> None:\n",
    "        print(delta.value, end=\"\", flush=True)\n",
    "\n",
    "    @override\n",
    "    async def on_run_step_created(self, run_step: RunStep) -> None:\n",
    "        details = run_step.step_details\n",
    "        if details.type == \"tool_calls\":\n",
    "            for tool in details.tool_calls:\n",
    "                if tool.type == \"code_interpreter\":\n",
    "                    print(\"\\nGenerating code to interpret:\\n\\n```python\")\n",
    "\n",
    "    @override\n",
    "    async def on_run_step_done(self, run_step: RunStep) -> None:\n",
    "        details = run_step.step_details\n",
    "        if details.type == \"tool_calls\":\n",
    "            for tool in details.tool_calls:\n",
    "                if tool.type == \"code_interpreter\":\n",
    "                    print(\"\\n```\\nExecuting code...\")\n",
    "\n",
    "    @override\n",
    "    async def on_run_step_delta(self, delta: RunStepDelta, snapshot: RunStep) -> None:\n",
    "        details = delta.step_details\n",
    "        if details is not None and details.type == \"tool_calls\":\n",
    "            for tool in details.tool_calls or []:\n",
    "                if tool.type == \"code_interpreter\" and tool.code_interpreter and tool.code_interpreter.input:\n",
    "                    print(tool.code_interpreter.input, end=\"\", flush=True)\n",
    "\n",
    "    @override\n",
    "    async def on_message_created(self, message: Message) -> None:\n",
    "        print(f\"{'-'*80}\\nAssistant:\\n\")\n",
    "\n",
    "    @override\n",
    "    async def on_message_done(self, message: Message) -> None:\n",
    "        # print a citation to the file searched\n",
    "        if not message.content:\n",
    "            return\n",
    "        content = message.content[0]\n",
    "        if not content.type == \"text\":\n",
    "            return\n",
    "        text_content = content.text\n",
    "        annotations = text_content.annotations\n",
    "        citations: List[str] = []\n",
    "        for index, annotation in enumerate(annotations):\n",
    "            text_content.value = text_content.value.replace(annotation.text, f\"[{index}]\")\n",
    "            if file_citation := getattr(annotation, \"file_citation\", None):\n",
    "                client = AsyncClient()\n",
    "                cited_file = await client.files.retrieve(file_citation.file_id)\n",
    "                citations.append(f\"[{index}] {cited_file.filename}\")\n",
    "        if citations:\n",
    "            print(\"\\n\".join(citations))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using the Agent\n",
    "\n",
    "First we need to use the `openai` client to create the actual assistant,\n",
    "thread, and vector store. Our AutoGen agent will be using these."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "\n",
    "# Create an assistant with code interpreter and file search tools.\n",
    "oai_assistant = openai.beta.assistants.create(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    description=\"An AI assistant that helps with everyday tasks.\",\n",
    "    instructions=\"Help the user with their task.\",\n",
    "    tools=[{\"type\": \"code_interpreter\"}, {\"type\": \"file_search\"}],\n",
    ")\n",
    "\n",
    "# Create a vector store to be used for file search.\n",
    "vector_store = openai.vector_stores.create()\n",
    "\n",
    "# Create a thread which is used as the memory for the assistant.\n",
    "thread = openai.beta.threads.create(\n",
    "    tool_resources={\"file_search\": {\"vector_store_ids\": [vector_store.id]}},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then, we create a runtime, and register an agent factory function for this \n",
    "agent with the runtime."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_core import SingleThreadedAgentRuntime\n",
    "\n",
    "runtime = SingleThreadedAgentRuntime()\n",
    "await OpenAIAssistantAgent.register(\n",
    "    runtime,\n",
    "    \"assistant\",\n",
    "    lambda: OpenAIAssistantAgent(\n",
    "        description=\"OpenAI Assistant Agent\",\n",
    "        client=openai.AsyncClient(),\n",
    "        assistant_id=oai_assistant.id,\n",
    "        thread_id=thread.id,\n",
    "        assistant_event_handler_factory=lambda: EventHandler(),\n",
    "    ),\n",
    ")\n",
    "agent = AgentId(\"assistant\", \"default\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's turn on logging to see what's happening under the hood."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "logging.basicConfig(level=logging.WARNING)\n",
    "logging.getLogger(\"autogen_core\").setLevel(logging.DEBUG)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's send a greeting message to the agent, and see the response streamed back."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Sending message of type TextMessage to assistant: {'content': 'Hello, how are you today!', 'source': 'user'}\n",
      "INFO:autogen_core:Calling message handler for assistant:default with message type TextMessage sent by Unknown\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "Assistant:\n",
      "\n",
      "Hello! I'm here and ready to assist you. How can I help you today?"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Resolving response with message type TextMessage for recipient None from assistant: {'content': \"Hello! I'm here and ready to assist you. How can I help you today?\", 'source': 'assistant'}\n"
     ]
    }
   ],
   "source": [
    "runtime.start()\n",
    "await runtime.send_message(TextMessage(content=\"Hello, how are you today!\", source=\"user\"), agent)\n",
    "await runtime.stop_when_idle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assistant with Code Interpreter\n",
    "\n",
    "Let's ask some math question to the agent, and see it uses the code interpreter\n",
    "to answer the question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Sending message of type TextMessage to assistant: {'content': 'What is 1332322 x 123212?', 'source': 'user'}\n",
      "INFO:autogen_core:Calling message handler for assistant:default with message type TextMessage sent by Unknown\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Calculating the product of 1332322 and 123212\n",
      "result = 1332322 * 123212\n",
      "result\n",
      "```\n",
      "Executing code...\n",
      "--------------------------------------------------------------------------------\n",
      "Assistant:\n",
      "\n",
      "The product of 1,332,322 and 123,212 is 164,158,058,264."
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Resolving response with message type TextMessage for recipient None from assistant: {'content': 'The product of 1,332,322 and 123,212 is 164,158,058,264.', 'source': 'assistant'}\n"
     ]
    }
   ],
   "source": [
    "runtime.start()\n",
    "await runtime.send_message(TextMessage(content=\"What is 1332322 x 123212?\", source=\"user\"), agent)\n",
    "await runtime.stop_when_idle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get some data from Seattle Open Data portal. We will be using the\n",
    "[City of Seattle Wage Data](https://data.seattle.gov/City-Business/City-of-Seattle-Wage-Data/2khk-5ukd/).\n",
    "Let's download it first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "response = requests.get(\"https://data.seattle.gov/resource/2khk-5ukd.csv\")\n",
    "with open(\"seattle_city_wages.csv\", \"wb\") as file:\n",
    "    file.write(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's send the file to the agent using an `UploadForCodeInterpreter` message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Sending message of type UploadForCodeInterpreter to assistant: {'file_path': 'seattle_city_wages.csv'}\n",
      "INFO:autogen_core:Calling message handler for assistant:default with message type UploadForCodeInterpreter sent by Unknown\n",
      "INFO:autogen_core:Resolving response with message type NoneType for recipient None from assistant: None\n"
     ]
    }
   ],
   "source": [
    "runtime.start()\n",
    "await runtime.send_message(UploadForCodeInterpreter(file_path=\"seattle_city_wages.csv\"), agent)\n",
    "await runtime.stop_when_idle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now ask some questions about the data to the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Sending message of type TextMessage to assistant: {'content': 'Take a look at the uploaded CSV file.', 'source': 'user'}\n",
      "INFO:autogen_core:Calling message handler for assistant:default with message type TextMessage sent by Unknown\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "import pandas as pd\n",
      "\n",
      "# Load the uploaded CSV file to examine its contents\n",
      "file_path = '/mnt/data/file-oEvRiyGyHc2jZViKyDqL8aoh'\n",
      "csv_data = pd.read_csv(file_path)\n",
      "\n",
      "# Display the first few rows of the dataframe to understand its structure\n",
      "csv_data.head()\n",
      "```\n",
      "Executing code...\n",
      "--------------------------------------------------------------------------------\n",
      "Assistant:\n",
      "\n",
      "The uploaded CSV file contains the following columns:\n",
      "\n",
      "1. **department**: The department in which the individual works.\n",
      "2. **last_name**: The last name of the employee.\n",
      "3. **first_name**: The first name of the employee.\n",
      "4. **job_title**: The job title of the employee.\n",
      "5. **hourly_rate**: The hourly rate for the employee's position.\n",
      "\n",
      "Here are the first few entries from the file:\n",
      "\n",
      "| department                     | last_name | first_name | job_title                          | hourly_rate |\n",
      "|--------------------------------|-----------|------------|------------------------------------|-------------|\n",
      "| Police Department              | Aagard    | Lori       | Pol Capt-Precinct                 | 112.70      |\n",
      "| Police Department              | Aakervik  | Dag        | Pol Ofcr-Detective                | 75.61       |\n",
      "| Seattle City Light             | Aaltonen  | Evan       | Pwrline Clear Tree Trimmer        | 53.06       |\n",
      "| Seattle Public Utilities       | Aar       | Abdimallik | Civil Engrng Spec,Sr               | 64.43       |\n",
      "| Seattle Dept of Transportation | Abad      | Abigail    | Admin Spec II-BU                  | 37.40       |\n",
      "\n",
      "If you need any specific analysis or information from this data, please let me know!"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Resolving response with message type TextMessage for recipient None from assistant: {'content': \"The uploaded CSV file contains the following columns:\\n\\n1. **department**: The department in which the individual works.\\n2. **last_name**: The last name of the employee.\\n3. **first_name**: The first name of the employee.\\n4. **job_title**: The job title of the employee.\\n5. **hourly_rate**: The hourly rate for the employee's position.\\n\\nHere are the first few entries from the file:\\n\\n| department                     | last_name | first_name | job_title                          | hourly_rate |\\n|--------------------------------|-----------|------------|------------------------------------|-------------|\\n| Police Department              | Aagard    | Lori       | Pol Capt-Precinct                 | 112.70      |\\n| Police Department              | Aakervik  | Dag        | Pol Ofcr-Detective                | 75.61       |\\n| Seattle City Light             | Aaltonen  | Evan       | Pwrline Clear Tree Trimmer        | 53.06       |\\n| Seattle Public Utilities       | Aar       | Abdimallik | Civil Engrng Spec,Sr               | 64.43       |\\n| Seattle Dept of Transportation | Abad      | Abigail    | Admin Spec II-BU                  | 37.40       |\\n\\nIf you need any specific analysis or information from this data, please let me know!\", 'source': 'assistant'}\n"
     ]
    }
   ],
   "source": [
    "runtime.start()\n",
    "await runtime.send_message(TextMessage(content=\"Take a look at the uploaded CSV file.\", source=\"user\"), agent)\n",
    "await runtime.stop_when_idle()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Sending message of type TextMessage to assistant: {'content': 'What are the top-10 salaries?', 'source': 'user'}\n",
      "INFO:autogen_core:Calling message handler for assistant:default with message type TextMessage sent by Unknown\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Sorting the data by hourly_rate in descending order and selecting the top 10 salaries\n",
      "top_10_salaries = csv_data[['first_name', 'last_name', 'job_title', 'hourly_rate']].sort_values(by='hourly_rate', ascending=False).head(10)\n",
      "top_10_salaries.reset_index(drop=True, inplace=True)\n",
      "top_10_salaries\n",
      "```\n",
      "Executing code...\n",
      "--------------------------------------------------------------------------------\n",
      "Assistant:\n",
      "\n",
      "Here are the top 10 salaries based on the hourly rates from the CSV file:\n",
      "\n",
      "| First Name | Last Name | Job Title                          | Hourly Rate |\n",
      "|------------|-----------|------------------------------------|-------------|\n",
      "| Eric       | Barden    | Executive4                        | 139.61      |\n",
      "| Idris      | Beauregard| Executive3                        | 115.90      |\n",
      "| Lori       | Aagard    | Pol Capt-Precinct                 | 112.70      |\n",
      "| Krista     | Bair      | Pol Capt-Precinct                 | 108.74      |\n",
      "| Amy        | Bannister | Fire Chief, Dep Adm-80 Hrs        | 104.07      |\n",
      "| Ginger     | Armbruster| Executive2                        | 102.42      |\n",
      "| William    | Andersen  | Executive2                        | 102.42      |\n",
      "| Valarie    | Anderson  | Executive2                        | 102.42      |\n",
      "| Paige      | Alderete  | Executive2                        | 102.42      |\n",
      "| Kathryn    | Aisenberg | Executive2                        | 100.65      |\n",
      "\n",
      "If you need any further details or analysis, let me know!"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Resolving response with message type TextMessage for recipient None from assistant: {'content': 'Here are the top 10 salaries based on the hourly rates from the CSV file:\\n\\n| First Name | Last Name | Job Title                          | Hourly Rate |\\n|------------|-----------|------------------------------------|-------------|\\n| Eric       | Barden    | Executive4                        | 139.61      |\\n| Idris      | Beauregard| Executive3                        | 115.90      |\\n| Lori       | Aagard    | Pol Capt-Precinct                 | 112.70      |\\n| Krista     | Bair      | Pol Capt-Precinct                 | 108.74      |\\n| Amy        | Bannister | Fire Chief, Dep Adm-80 Hrs        | 104.07      |\\n| Ginger     | Armbruster| Executive2                        | 102.42      |\\n| William    | Andersen  | Executive2                        | 102.42      |\\n| Valarie    | Anderson  | Executive2                        | 102.42      |\\n| Paige      | Alderete  | Executive2                        | 102.42      |\\n| Kathryn    | Aisenberg | Executive2                        | 100.65      |\\n\\nIf you need any further details or analysis, let me know!', 'source': 'assistant'}\n"
     ]
    }
   ],
   "source": [
    "runtime.start()\n",
    "await runtime.send_message(TextMessage(content=\"What are the top-10 salaries?\", source=\"user\"), agent)\n",
    "await runtime.stop_when_idle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Assistant with File Search\n",
    "\n",
    "Let's try the Q&A over document feature. We first download Wikipedia page\n",
    "on the Third Anglo-Afghan War."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = requests.get(\"https://en.wikipedia.org/wiki/Third_Anglo-Afghan_War\")\n",
    "with open(\"third_anglo_afghan_war.html\", \"wb\") as file:\n",
    "    file.write(response.content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Send the file to the agent using an `UploadForFileSearch` message."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Sending message of type UploadForFileSearch to assistant: {'file_path': 'third_anglo_afghan_war.html', 'vector_store_id': 'vs_h3xxPbJFnd1iZ9WdjsQwNdrp'}\n",
      "INFO:autogen_core:Calling message handler for assistant:default with message type UploadForFileSearch sent by Unknown\n",
      "INFO:autogen_core:Resolving response with message type NoneType for recipient None from assistant: None\n"
     ]
    }
   ],
   "source": [
    "runtime.start()\n",
    "await runtime.send_message(\n",
    "    UploadForFileSearch(file_path=\"third_anglo_afghan_war.html\", vector_store_id=vector_store.id), agent\n",
    ")\n",
    "await runtime.stop_when_idle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's ask some questions about the document to the agent. Before asking,\n",
    "we reset the agent memory to start a new conversation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Sending message of type Reset to assistant: {}\n",
      "INFO:autogen_core:Calling message handler for assistant:default with message type Reset sent by Unknown\n",
      "INFO:autogen_core:Resolving response with message type NoneType for recipient None from assistant: None\n",
      "INFO:autogen_core:Sending message of type TextMessage to assistant: {'content': 'When and where was the treaty of Rawalpindi signed? Answer using the document provided.', 'source': 'user'}\n",
      "INFO:autogen_core:Calling message handler for assistant:default with message type TextMessage sent by Unknown\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "--------------------------------------------------------------------------------\n",
      "Assistant:\n",
      "\n",
      "The Treaty of Rawalpindi was signed on **8 August 1919**. The location of the signing was in **Rawalpindi**, which is in present-day Pakistan【6:0†source】."
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:autogen_core:Resolving response with message type TextMessage for recipient None from assistant: {'content': 'The Treaty of Rawalpindi was signed on **8 August 1919**. The location of the signing was in **Rawalpindi**, which is in present-day Pakistan【6:0†source】.', 'source': 'assistant'}\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0] third_anglo_afghan_war.html\n"
     ]
    }
   ],
   "source": [
    "runtime.start()\n",
    "await runtime.send_message(Reset(), agent)\n",
    "await runtime.send_message(\n",
    "    TextMessage(\n",
    "        content=\"When and where was the treaty of Rawalpindi signed? Answer using the document provided.\", source=\"user\"\n",
    "    ),\n",
    "    agent,\n",
    ")\n",
    "await runtime.stop_when_idle()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "That's it! We have successfully built an agent backed by OpenAI Assistant."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
